import pandas as pd
import torch
from torch.utils.data import DataLoader
import os
from prepare_data import DACDataset
from train.deep_fm import DeepFM
from sklearn.preprocessing import LabelEncoder, MinMaxScaler


def get_predictions(model, test_dataset, batch_size=1024):
    # 1. 切换模型为推理模式（关闭Dropout等训练层）
    model.eval()

    # 2. 创建测试集数据加载器（按批次加载，避免内存溢出）
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,  # 测试集无需打乱
        num_workers=0  # 单线程加载，避免多进程问题
    )

    # 3. 存储所有预测结果
    predictions = []

    # 4. 关闭梯度计算，加速推理
    with torch.no_grad():
        for batch in test_loader:
            # 从数据加载器获取批次数据（根据DACDataset的__getitem__返回格式调整）
            # 假设每个batch包含类别特征和数值特征，格式为 (cat_feats, num_feats)
            cat_feats, num_feats = batch

            # 模型推理：输入类别特征和数值特征，输出点击率预测概率
            outputs = model(cat_feats, num_feats)

            # 将预测结果转换为numpy数组并添加到列表（flatten()确保为一维）
            predictions.extend(outputs.numpy().flatten())

    return predictions


def main():
    sample_size = 10000
    batch_size = 1000
    embedding_dim = 8
    hidden_dims = [512, 256, 128, 64]
    # 添加信任的类到安全列表
    torch.serialization.add_safe_globals([
        LabelEncoder,
        MinMaxScaler,
        # 若还有其他类报错，按提示添加到这里
    ])
    test_path = '../kaggle-display-advertising-challenge-dataset/test.txt'
    if not os.path.exists(test_path):
        raise FileNotFoundError(f"数据文件不存在: {test_path}")

    # 加载保存的文件
    checkpoint = torch.load('../res/deepfm_dac_v1_final.pth', map_location='cpu', weights_only=False)

    # 提取预处理工具
    cat_encoders = checkpoint['cat_encoders']  # 每个类别特征的LabelEncoder
    num_scaler = checkpoint['num_scaler']
    cat_dims = [len(encoder.classes_) for encoder in cat_encoders.values()]
    model = DeepFM(
        cat_dims=cat_dims,
        num_cols_count=13,
        embedding_dim=embedding_dim,
        hidden_dims=hidden_dims
    )
    model.load_state_dict(checkpoint['model_state_dict'])

    test_dataset = DACDataset(
        file_path=test_path,
        cat_encoders=cat_encoders,
        num_scaler=num_scaler,
        sample_size=sample_size
    )



    # 调用函数获取预测结果
    test_predictions = get_predictions(model, test_dataset, batch_size)

    # 5. 保存预测结果（按Kaggle比赛要求格式，通常为无表头的单列概率）
    pd.DataFrame(test_predictions).to_csv(
        '../res/deepfm_test_predictions.csv',
        index=False,  # 不保存索引
        header=False  # 不保存表头
    )

    print(f"预测完成，共生成{len(test_predictions)}条结果，已保存至deepfm_test_predictions.csv")


if __name__ == "__main__":
    main()
